import pandas as pd
raw = pd.read_csv("金庸-射雕英雄传txt精校版.txt",
                  names=['txt'],
                  sep='aaa',
                  encoding="utf-8",
                  engine='python')
print(len(raw))


def m_head(tmpstr):
    return tmpstr[:1]


def m_mid(tmpstr):
    return tmpstr.find("回 ")


raw['head'] = raw.txt.apply(m_head)
raw['mid'] = raw.txt.apply(m_mid)
raw['len'] = raw.txt.apply(len)
# raw['chap'] = 0
raw.head(50)
# 章节判断
chapnum = 0
for i in range(len(raw)):
    if raw['head'][i] == "第" and raw['mid'][i] > 0 and raw['len'][i] < 30:
        chapnum += 1
    if chapnum >= 40 and raw['txt'][i] == "附录一：成吉思汗家族":
        chapnum = 0
    raw.loc[i, 'chap'] = chapnum

# 删除临时变量
del raw['head']
del raw['mid']
del raw['len']
raw.head(50)

stoplist = list(
    pd.read_csv('停用词.txt',
                names=['w'],
                sep='aaa',
                encoding='utf-8',
                engine='python').w)
import jieba


def m_cut(intxt):
    return [w for w in jieba.cut(intxt) if w not in stoplist and len(w) > 1]


chaplist = []
for i in set(raw.chap):
    tmpchap = raw[raw.chap == i]
    s = [m_cut(w) for w in tmpchap.txt]
    chaplist.extend(s)

from gensim import corpora, models

# 生成文档对应的字典和bow稀疏向量

dictionary = corpora.Dictionary(chaplist)
corpus = [dictionary.doc2bow(text) for text in chaplist]  # 仍为list in list
tfidf_model = models.TfidfModel(corpus)  # 建立TF-IDF模型
corpus_tfidf = tfidf_model[corpus]  # 对所需文档计算TF-IDF结果

from gensim.models.ldamodel import LdaModel

# 列出所消耗的时间备查
ldamodel = LdaModel(corpus, id2word=dictionary, num_topics=10, passes=2)
ldamodel20 = LdaModel(corpus, id2word=dictionary, num_topics=20, passes=2)
ldamodel30 = LdaModel(corpus, id2word=dictionary, num_topics=30, passes=2)

topic1 = ldamodel.print_topics()
print("10个主题")
for i in topic1:
    print(i)
topic2 = ldamodel20.print_topics()
print("20个主题")
for i in topic2:
    print(i)
topic3 = ldamodel30.print_topics()
print("30个主题")
for i in topic3:
    print(i)

for i in set(raw.chap):
    if i % 2 == 1:
        topic = ldamodel.get_document_topics(corpus[int(i)])
        topic = max(topic, key=lambda x: x[1])
        print('第', str(i), '章在10个主题的编号为', topic)

for i in set(raw.chap):
    if i % 2 == 1:
        topic = ldamodel.get_document_topics(corpus[int(i)])
        topic = max(topic, key=lambda x: x[1])
        print('第', str(i), '章在20个主题的编号', topic)

for i in set(raw.chap):
    if i % 2 == 1:
        topic = ldamodel.get_document_topics(corpus[int(i)])
        topic = max(topic, key=lambda x: x[1])
        print('第', str(i), '章在30个主题的编号', topic)
